package resolver

import (
	"fmt"
	"github.com/miekg/dns"
	"net"
	"strings"
	"time"
)

func Runner(c *dns.Client, m *dns.Msg, dnsServer string) (*DomainDNSRecord, *dns.Msg, error, bool) {
	d := NewDomainDNSRecord()
	flag := false
	// 进行DNS查询
	r, rtt, err := c.Exchange(m, dnsServer)
	// fmt.Println(domain, "查询耗时：", rtt)
	if err != nil {
		if strings.HasSuffix(strings.TrimSpace(err.Error()), "no suitable address found") {
			res1 := strings.Split(err.Error(), "[")
			if len(res1) > 1 {
				res2 := strings.Split(res1[1], "]")
				if len(res2) > 1 {
					// fmt.Printf("源地址 <%s> 无法向指定的DNS <%s> 发起请求\n", res2[0], dnsServer)
					return nil, nil, fmt.Errorf("源地址 <%s> 无法向指定的DNS <%s> 发起请求", res2[0], dnsServer), false
				} else {
					// fmt.Printf("无法向指定的DNS <%s> 发起请求\n", dnsServer)
					return nil, nil, fmt.Errorf("无法向指定的DNS <%s> 发起请求", dnsServer), false
				}
			} else {
				// fmt.Printf("无法向指定的DNS <%s> 发起请求\n", dnsServer)
				return nil, nil, fmt.Errorf("无法向指定的DNS <%s> 发起请求", dnsServer), false
			}
		} else {
			if strings.HasSuffix(
				strings.TrimSpace(err.Error()),
				"A socket operation was attempted to an unreachable network.",
			) {
				// fmt.Println("网络不可达")
				return nil, nil, fmt.Errorf("网络不可达"), false
			} else if strings.HasSuffix(
				strings.TrimSpace(err.Error()),
				"i/o timeout",
			) {
				// fmt.Printf("DNS %s 请求超时\n", dnsServer)
				return nil, nil, fmt.Errorf("DNS %s 请求超时", dnsServer), false
			} else if strings.HasSuffix(
				strings.TrimSpace(err.Error()),
				"no such host",
			) {
				// fmt.Printf("DNS %s 不合法\n", dnsServer)
				return nil, nil, fmt.Errorf("DNS %s 不合法", dnsServer), false
			} else {
				// fmt.Printf("查询出错: %v\n", err)
				return nil, nil, fmt.Errorf("查询出错: %v", err), false
			}
		}
	} else {
		if r.Rcode == dns.RcodeSuccess {
			flag = true
		}
		d.SetAuthoritative(r.Authoritative)
		d.SetRecursionAvailable(r.RecursionAvailable)
		d.SetDNSServer(dnsServer)
		d.SetDuration(rtt)
		return d, r, nil, flag
	}
}

// ResolveDNS 解析域名
//
// 参数：
//
// domain: 待解析的域名
// recordType: 待解析的记录类型
// server: DNS服务器地址
// port: DNS服务器端口
// provider: 开启EDNS0 Client Subnet时需要提供的ip或者网段
// localIP: 需要绑定源地址时提供的本地IP地址
// localPort: 需要绑定源地址时提供的本地端口
//
// 返回值：
//
//	DomainDNSRecord: 域名解析记录详情
//	error: 错误信息
//	bool: 是否解析成功
func ResolveDNS(domain string, recordType RecordType, server string, port int, provider, localIP string, localPort int) (*DomainDNSRecord, error, bool) {
	ip := net.ParseIP(server)
	dnsServer := ""
	if ip == nil {
		dnsServer = fmt.Sprintf("%s:%d", server, port)
	} else {
		if ip.To4() != nil {
			dnsServer = fmt.Sprintf("%s:%d", server, port)
		} else {
			dnsServer = fmt.Sprintf("[%s]:%d", server, port)
		}
	}
	var dialer *net.Dialer
	if localIP != "" {
		if localPort != 0 {
			dialer = &net.Dialer{
				LocalAddr: &net.UDPAddr{
					IP:   net.ParseIP(localIP),
					Port: localPort,
				},
			}
		} else {
			// 创建一个DNS客户端，并设置源地址
			dialer = &net.Dialer{
				LocalAddr: &net.UDPAddr{
					IP: net.ParseIP(localIP),
				},
			}
		}
	}
	c := new(dns.Client)
	c.Timeout = 5 * time.Second
	if dialer != nil {
		c.Dialer = dialer
	}
	// 创建一个DNS查询消息
	m := new(dns.Msg)
	// 允许递归查询
	m.RecursionDesired = true

	if provider != "" {
		opt := new(dns.OPT)
		opt.Hdr.Name = "."
		opt.Hdr.Rrtype = dns.TypeOPT
		ip := net.ParseIP(provider)
		if ip == nil {
			return nil, fmt.Errorf("无效的provider: %s", provider), false
		}
		ecs := new(dns.EDNS0_SUBNET)
		if ip.To4() != nil {
			ecs = &dns.EDNS0_SUBNET{
				Code:          dns.EDNS0SUBNET,
				Family:        1, // 1 表示 IPv4，2 表示 IPv6
				SourceNetmask: 24,
				SourceScope:   0,
				Address:       ip.To4(),
			}
		} else {
			ecs = &dns.EDNS0_SUBNET{
				Code:          dns.EDNS0SUBNET,
				Family:        2, // 1 表示 IPv4，2 表示 IPv6
				SourceNetmask: 56,
				SourceScope:   0,
				Address:       ip.To16(),
			}
		}

		// 创建EDNS0客户端子网选项
		opt.Option = append(opt.Option, ecs)
		m.Extra = append(m.Extra, opt)
	}
	switch recordType {
	case AAAA:
		m.SetQuestion(dns.Fqdn(domain), dns.TypeAAAA)
		if d, r, err, ok := Runner(c, m, dnsServer); ok {
			// 处理响应
			for _, ans := range r.Answer {
				if a, ok := ans.(*dns.AAAA); ok {
					// fmt.Printf("AAAA Record: %s -> %s\n", domain, a.AAAA.String())
					d.AddResult(a.AAAA.String())
					d.AddFullResult(a.String())
				}
			}
			return d, err, ok
		}
	case NS:
		m.SetQuestion(dns.Fqdn(domain), dns.TypeNS)
		if d, r, err, ok := Runner(c, m, dnsServer); ok {
			// 处理响应
			for _, ans := range r.Answer {
				if a, ok := ans.(*dns.NS); ok {
					d.AddResult(a.Ns)
					d.AddFullResult(a.String())
				}
			}
			return d, err, ok
		}
	case CNAME:
		m.SetQuestion(dns.Fqdn(domain), dns.TypeCNAME)
		if d, r, err, ok := Runner(c, m, dnsServer); ok {
			// 处理响应
			for _, ans := range r.Answer {
				if a, ok := ans.(*dns.CNAME); ok {
					d.AddResult(a.Target)
					d.AddFullResult(a.String())
				}
			}
			return d, err, ok
		}
	case MX:
		m.SetQuestion(dns.Fqdn(domain), dns.TypeMX)
		if d, r, err, ok := Runner(c, m, dnsServer); ok {
			// 处理响应
			for _, ans := range r.Answer {
				if a, ok := ans.(*dns.MX); ok {
					d.AddResult(a.Mx)
					d.AddFullResult(a.String())
				}
			}
			return d, err, ok
		}
	case TXT:
		m.SetQuestion(dns.Fqdn(domain), dns.TypeTXT)
		if d, r, err, ok := Runner(c, m, dnsServer); ok {
			// 处理响应
			for _, ans := range r.Answer {
				if a, ok := ans.(*dns.TXT); ok {
					for _, txt := range a.Txt {
						d.AddResult(txt)
						d.AddFullResult(a.String())
					}
				}
			}
			return d, err, ok
		}
	default:
		m.SetQuestion(dns.Fqdn(domain), dns.TypeA)
		if d, r, err, ok := Runner(c, m, dnsServer); ok {
			// 处理响应
			for _, ans := range r.Answer {
				if a, ok := ans.(*dns.A); ok {
					d.AddResult(a.A.String())
					d.AddFullResult(a.String())
				}
			}
			return d, err, ok
		} else {
			return nil, err, false
		}
	}
	return nil, nil, false
}

// SimpleResolveDNS 简单解析，默认端口为53
//
// 不绑定本地地址端口
// dns server端口为53
func SimpleResolveDNS(domain string, recordType RecordType, server string) (*DomainDNSRecord, error, bool) {
	return ResolveDNS(domain, recordType, server, 53, "", "", 0)
}

func SimpleAndPortResolveDNS(domain string, recordType RecordType, server string, port int) (*DomainDNSRecord, error, bool) {
	return ResolveDNS(domain, recordType, server, port, "", "", 0)
}

func SimpleECSResolveDNS(domain string, recordType RecordType, server string, subNet string) (*DomainDNSRecord, error, bool) {
	return ResolveDNS(domain, recordType, server, 53, subNet, "", 0)
}

func SimpleECSAndPortResolveDNS(domain string, recordType RecordType, server string, port int, subNet string) (*DomainDNSRecord, error, bool) {
	return ResolveDNS(domain, recordType, server, port, subNet, "", 0)
}

func SimpleAndSourceIPResolveDNS(domain string, recordType RecordType, server string, localIP string, localPort int) (*DomainDNSRecord, error, bool) {
	return ResolveDNS(domain, recordType, server, 53, "", localIP, localPort)
}

// MultipleServerResolveDNS 多dns服务器解析，解析成功就不再解析后面的dns服务器
//
// 参数：
//
//	domain: 待解析的域名
//	recordType: 待解析的记录类型
//	servers: 多个DNS服务器地址
//	port: DNS服务器端口
//	subNet: 开启EDNS0 Client Subnet时需要提供的ip或者网段
//	localIP: 需要绑定源地址时提供的本地IP地址
//	localPort: 需要绑定源地址时提供的本地端口
//
// 返回值：
//
//	DomainDNSRecord: 域名解析记录详情
//	error: 错误信息
func MultipleServerResolveDNS(domain string, recordType RecordType, servers []string, port, localPort int, localIP, subNet string) (*DomainDNSRecord, error) {
	for _, server := range servers {
		res, _, ok := ResolveDNS(domain, recordType, server, port, subNet, localIP, localPort)
		if ok {
			return res, nil
		}
	}
	return nil, fmt.Errorf("no result")
}

// 快捷方式

func SimpleMultipleServerResolveDNS(domain string, recordType RecordType, servers []string) (*DomainDNSRecord, error) {
	return MultipleServerResolveDNS(domain, recordType, servers, 53, 0, "", "")
}
func SimpleMultipleServerResolveDNSWithPort(domain string, recordType RecordType, servers []string, port int) (*DomainDNSRecord, error) {
	return MultipleServerResolveDNS(domain, recordType, servers, port, 0, "", "")
}

func SimpleMultipleServerResolveDNSWithECS(domain string, recordType RecordType, servers []string, subNet string) (*DomainDNSRecord, error) {
	return MultipleServerResolveDNS(domain, recordType, servers, 53, 0, "", subNet)
}
func SimpleMultipleServerResolveDNSWithECSAndPort(domain string, recordType RecordType, servers []string, port int, subNet string) (*DomainDNSRecord, error) {
	return MultipleServerResolveDNS(domain, recordType, servers, port, 0, "", subNet)
}
func SimpleMultipleServerResolveDNSWithSourceIP(domain string, recordType RecordType, servers []string, localIP string, localPort int) (*DomainDNSRecord, error) {
	return MultipleServerResolveDNS(domain, recordType, servers, 53, localPort, localIP, "")
}

// OneServerResolveDNS 单dns服务器解析
//
// 参数：
//
//	domain: 待解析的域名
//	recordType: 待解析的记录类型
//	server: DNS服务器地址
//	port: DNS服务器端口
//	subNet: 开启EDNS0 Client Subnet时需要提供的ip或者网段
//	localIP: 需要绑定源地址时提供的本地IP地址
//	localPort: 需要绑定源地址时提供的本地端口
//
// 返回值：
//
//	DomainDNSRecord: 域名解析记录详情
//	error: 错误信息
func OneServerResolveDNS(domain string, recordType RecordType, server string, port, localPort int, localIP, subNet string) (*DomainDNSRecord, error) {
	res, err, ok := ResolveDNS(domain, recordType, server, port, subNet, localIP, localPort)
	if ok {
		return res, nil
	}
	return nil, err
}

// SimpleOneServerResolveDNS 提供了一个默认端口为53，本地端口为0，本地IP和子网为空的简化版本。
func SimpleOneServerResolveDNS(domain string, recordType RecordType, server string) (*DomainDNSRecord, error) {
	return OneServerResolveDNS(domain, recordType, server, 53, 0, "", "")
}

// SimpleOneServerResolveDNSWithPort 提供了一个指定端口的简化版本。
func SimpleOneServerResolveDNSWithPort(domain string, recordType RecordType, server string, port int) (*DomainDNSRecord, error) {
	return OneServerResolveDNS(domain, recordType, server, port, 0, "", "")
}

// SimpleOneServerResolveDNSWithECS 提供了一个指定子网的简化版本。
func SimpleOneServerResolveDNSWithECS(domain string, recordType RecordType, server string, subNet string) (*DomainDNSRecord, error) {
	return OneServerResolveDNS(domain, recordType, server, 53, 0, "", subNet)
}

// SimpleOneServerResolveDNSWithECSAndPort 提供了一个指定端口和子网的简化版本。
func SimpleOneServerResolveDNSWithECSAndPort(domain string, recordType RecordType, server string, port int, subNet string) (*DomainDNSRecord, error) {
	return OneServerResolveDNS(domain, recordType, server, port, 0, "", subNet)
}

// SimpleOneServerResolveDNSWithSourceIP 提供了一个指定本地IP和本地端口的简化版本。
func SimpleOneServerResolveDNSWithSourceIP(domain string, recordType RecordType, server string, localIP string, localPort int) (*DomainDNSRecord, error) {
	return OneServerResolveDNS(domain, recordType, server, 53, localPort, localIP, "")
}

// PrintToResult 打印解析结果
//
// 参数：
//
// short: 是否只打印解析结果 简单的
// count: 解析次数
// domain: 待解析的域名
// recordType: 待解析的记录类型
// servers: 多DNS服务器地址 哪怕只有一个DNS server 也要写成切片形式
// port: DNS服务器端口
// subNet: 开启EDNS0 Client Subnet时需要提供的ip或者网段
// localIP: 需要绑定源地址时提供的本地IP地址
// localPort: 需要绑定源地址时提供的本地端口
// manual: 是否手动指定DNS服务器 如果为true，则意思是servers中的所有DNS服务器是手动的，需要全部解析出来；如果为false，则意思是servers中的DNS服务器是自动的，只要解析成功就返回。
func PrintToResult(short bool, count int, domain string, recordType RecordType, servers []string, port, localPort int, localIP, subNet string, manual bool) {
	if manual {
		for _, server := range servers {
			if len(servers) > 1 {
				fmt.Printf("======== DNS %s 的解析 %s 结果============\n", server, domain)
			}
			if short {
				res, err := OneServerResolveDNS(
					domain,
					recordType,
					server,
					port,
					localPort,
					localIP,
					subNet,
				)
				if err != nil {
					fmt.Println(err)
				} else {
					for _, v := range res.Result {
						fmt.Println(v)
					}
					fmt.Println("耗时:", res.Duration)
				}
			} else {
				for i := 1; i <= count; i++ {
					res, err := OneServerResolveDNS(
						domain,
						recordType,
						server,
						port,
						localPort,
						localIP,
						subNet,
					)
					if err == nil {
						if len(res.Result) == 0 {
							fmt.Printf("%-18s无 %s记录 解析结果\n", domain, recordType)
							fmt.Printf("%-14s%v\n", "支持递归:", res.RecursionAvailable)
						} else {
							if count > 1 {
								fmt.Printf("-------- %s 第 %d 次解析结果------------\n", domain, i)
							} else {
								fmt.Println(domain, "解析结果：")
							}
							for _, v := range res.FullResult {
								fmt.Println("+ ", v)
							}
						}
						fmt.Printf("%-18s%v\n", "SERVER:", res.DNSServer)
						fmt.Printf("%-11s%v\n", "来自权威服务器:", res.Authoritative)
						fmt.Printf("%-12s%v\n", "本次解析耗时:", res.Duration)
						if len(res.Result) == 0 {
							break
						}
					} else {
						fmt.Println(err)
					}
				}
			}
		}
	} else {
		if short {
			res, err := MultipleServerResolveDNS(domain, recordType, servers, port, localPort, localIP, subNet)
			if err != nil {
				fmt.Println(err)
				return
			}
			for _, v := range res.Result {
				fmt.Println(v)
			}
		} else {
			for i := 1; i <= count; i++ {
				res, err := MultipleServerResolveDNS(
					domain,
					recordType,
					servers,
					port,
					localPort,
					localIP,
					subNet,
				)
				if err != nil {
					fmt.Println(err)
					return
				}
				if len(res.Result) == 0 {
					fmt.Printf("%-18s无 %s记录 解析结果\n", domain, recordType)
					fmt.Printf("%-14s%v\n", "支持递归:", res.RecursionAvailable)
				} else {
					if count > 1 {
						fmt.Printf("-------- %s 第 %d 次解析结果------------\n", domain, i)
					} else {
						fmt.Println(domain, "解析结果：")
					}
					for _, v := range res.FullResult {
						fmt.Println("+ ", v)
					}
				}
				fmt.Printf("%-18s%v\n", "SERVER:", res.DNSServer)
				fmt.Printf("%-11s%v\n", "来自权威服务器:", res.Authoritative)
				fmt.Printf("%-12s%v\n", "本次解析耗时:", res.Duration)
				if len(res.Result) == 0 {
					break
				}
			}
		}
	}
}

// SimplePrintToResult 提供了一个默认端口为53，本地端口为0，本地IP和子网为空的简化版本。
func SimplePrintToResult(short bool, count int, domain string, recordType RecordType, servers []string) {
	PrintToResult(short, count, domain, recordType, servers, 53, 0, "", "", false)
}

// SimplePrintToResultWithShort 提供了一个简短输出的简化版本。
func SimplePrintToResultWithShort(count int, domain string, recordType RecordType, servers []string) {
	SimplePrintToResult(true, count, domain, recordType, servers)
}

// SimplePrintToResultWithDetail 提供了一个详细输出的简化版本。
func SimplePrintToResultWithDetail(count int, domain string, recordType RecordType, servers []string) {
	SimplePrintToResult(false, count, domain, recordType, servers)
}

// SimplePrintToResultWithPort 提供了一个指定端口的简化版本。
func SimplePrintToResultWithPort(short bool, count int, domain string, recordType RecordType, servers []string, port int) {
	PrintToResult(short, count, domain, recordType, servers, port, 0, "", "", false)
}

// SimplePrintToResultWithPortAndShort 提供了一个指定端口和简短输出的简化版本。
func SimplePrintToResultWithPortAndShort(count int, domain string, recordType RecordType, servers []string, port int) {
	SimplePrintToResultWithPort(true, count, domain, recordType, servers, port)
}

// SimplePrintToResultWithPortAndDetail 提供了一个指定端口和详细输出的简化版本。
func SimplePrintToResultWithPortAndDetail(count int, domain string, recordType RecordType, servers []string, port int) {
	SimplePrintToResultWithPort(false, count, domain, recordType, servers, port)
}

// SimplePrintToResultWithECS 提供了一个指定子网的简化版本。
func SimplePrintToResultWithECS(short bool, count int, domain string, recordType RecordType, servers []string, subNet string) {
	PrintToResult(short, count, domain, recordType, servers, 53, 0, "", subNet, false)
}

// SimplePrintToResultWithECSAndShort 提供了一个指定子网和简短输出的简化版本。
func SimplePrintToResultWithECSAndShort(count int, domain string, recordType RecordType, servers []string, subNet string) {
	SimplePrintToResultWithECS(true, count, domain, recordType, servers, subNet)
}

// SimplePrintToResultWithECSAndDetail 提供了一个指定子网和详细输出的简化版本。
func SimplePrintToResultWithECSAndDetail(count int, domain string, recordType RecordType, servers []string, subNet string) {
	SimplePrintToResultWithECS(false, count, domain, recordType, servers, subNet)
}

// SimplePrintToResultWithECSAndPort 提供了一个指定端口和子网的简化版本。
func SimplePrintToResultWithECSAndPort(short bool, count int, domain string, recordType RecordType, servers []string, port int, subNet string) {
	PrintToResult(short, count, domain, recordType, servers, port, 0, "", subNet, false)
}

// SimplePrintToResultWithECSAndPortAndShort 提供了一个指定端口、子网和简短输出的简化版本。
func SimplePrintToResultWithECSAndPortAndShort(count int, domain string, recordType RecordType, servers []string, port int, subNet string) {
	SimplePrintToResultWithECSAndPort(true, count, domain, recordType, servers, port, subNet)
}

// SimplePrintToResultWithECSAndPortAndDetail 提供了一个指定端口、子网和详细输出的简化版本。
func SimplePrintToResultWithECSAndPortAndDetail(count int, domain string, recordType RecordType, servers []string, port int, subNet string) {
	SimplePrintToResultWithECSAndPort(false, count, domain, recordType, servers, port, subNet)
}

// SimplePrintToResultWithSourceIP 提供了一个指定本地IP和本地端口的简化版本。
func SimplePrintToResultWithSourceIP(short bool, count int, domain string, recordType RecordType, servers []string, localIP string, localPort int) {
	PrintToResult(short, count, domain, recordType, servers, 53, localPort, localIP, "", false)
}

// SimplePrintToResultWithSourceIPAndShort 提供了一个指定本地IP、本地端口和简短输出的简化版本。
func SimplePrintToResultWithSourceIPAndShort(count int, domain string, recordType RecordType, servers []string, localIP string, localPort int) {
	SimplePrintToResultWithSourceIP(true, count, domain, recordType, servers, localIP, localPort)
}

// SimplePrintToResultWithSourceIPAndDetail 提供了一个指定本地IP、本地端口和详细输出的简化版本。
func SimplePrintToResultWithSourceIPAndDetail(count int, domain string, recordType RecordType, servers []string, localIP string, localPort int) {
	SimplePrintToResultWithSourceIP(false, count, domain, recordType, servers, localIP, localPort)
}

// GeneralPrintToResultWithPortAndShort 通用的快捷方式 打印解析结果
//
// 参数：
//
// domain: 待解析的域名
// recordType: 待解析的记录类型
// servers: 多DNS服务器地址 哪怕只有一个DNS server 也要写成切片形式
// port: DNS服务器端口
func GeneralPrintToResultWithPortAndShort(domain string, recordType RecordType, servers []string, port int) {
	// 这里默认不开启 EDNS0 Client Subnet，不绑定源IP，只解析一次
	SimplePrintToResultWithPortAndShort(1, domain, recordType, servers, port)
}

func GeneralPrintToResultWithPortAndDetail(domain string, recordType RecordType, servers []string, port int) {
	SimplePrintToResultWithPortAndDetail(1, domain, recordType, servers, port)
}

func GeneralPrintToResultAndShort(domain string, recordType RecordType, servers []string) {
	// 这里默认不开启 EDNS0 Client Subnet，不绑定源IP，只解析一次
	SimplePrintToResultWithPortAndShort(1, domain, recordType, servers, 53)
}

func GeneralPrintToResultAndDetail(domain string, recordType RecordType, servers []string) {
	SimplePrintToResultWithPortAndDetail(1, domain, recordType, servers, 53)
}

func GeneralManualPrintToResultWithPortAndShort(domain string, recordType RecordType, servers []string, port int) {
	PrintToResult(true, 1, domain, recordType, servers, port, 0, "", "", true)
}

func GeneralManualPrintToResultWithPortAndDetail(domain string, recordType RecordType, servers []string, port int) {
	PrintToResult(false, 1, domain, recordType, servers, port, 0, "", "", true)
}

func GeneralManualPrintToResultAndShort(domain string, recordType RecordType, servers []string) {
	// 这里默认不开启 EDNS0 Client Subnet，不绑定源IP，只解析一次
	PrintToResult(true, 1, domain, recordType, servers, 53, 0, "", "", true)
}

func GeneralManualPrintToResultAndDetail(domain string, recordType RecordType, servers []string) {
	PrintToResult(false, 1, domain, recordType, servers, 53, 0, "", "", true)
}

func GeneralPrintToResult(domain string, recordType RecordType, server string, port int, short bool) {
	PrintToResult(short, 1, domain, recordType, []string{server}, port, 0, "", "", false)
}
